package rabbit.flt.plugins.webflux.plugin;

import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpResponse;
import rabbit.flt.common.TraceContextHolder;
import rabbit.flt.common.context.TraceMetaData;
import rabbit.flt.common.trace.TraceData;
import rabbit.flt.common.trace.io.HttpResponse;
import rabbit.flt.plugins.common.plugin.SupportPlugin;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.lang.reflect.Method;
import java.nio.charset.Charset;

public class ServerHttpResponsePlugin extends SupportPlugin {

    @Override
    public Object[] before(Object target, Method method, Object[] args) {
        if (method.getName().equals("setComplete")) {
            fillTraceResponse(target);
            return args;
        }
        if (args[0] instanceof Mono) {
            args[0] = ((Mono<DataBuffer>) args[0]).map(buffer -> handleResponse(target, buffer));
        } else {
            args[0] = ((Flux<Object>) args[0]).map(source -> {
                if (source instanceof DataBuffer) {
                    return handleResponse(target, (DataBuffer) source);
                }
                if (source instanceof Publisher) {
                    Publisher<DataBuffer> publisher = (Publisher<DataBuffer>) source;
                    return Flux.from(publisher).map(buffer -> handleResponse(target, buffer));
                } else {
                    logger.error("unknown source type: {}", source.getClass().getName());
                    return source;
                }
            });
        }
        return args;
    }

    private TraceData fillTraceResponse(Object target) {
        ServerHttpResponse response = (ServerHttpResponse) target;
        TraceContextHolder holder = (TraceContextHolder) target;
        TraceMetaData metaData = (TraceMetaData) holder.getTraceContextData();
        TraceData traceData = metaData.getWebTraceData();
        if (null != traceData) {
            if (null == traceData.getHttpResponse()) {
                traceData.setHttpResponse(new HttpResponse());
            }
            HttpResponse httpResponse = traceData.getHttpResponse();
            httpResponse.setStatusCode(response.getRawStatusCode());
            response.getHeaders().forEach((name, value) -> httpResponse.addHeader(name, truncate(value.toString())));
            traceData.updateCost();
        }
        return traceData;
    }

    private DataBuffer handleResponse(Object target, DataBuffer buffer) {
        TraceData traceData = fillTraceResponse(target);
        if (null != traceData) {
            ServerHttpResponse response = (ServerHttpResponse) target;
            if (HttpStatus.OK != response.getStatusCode()) {
                // 非200请求记录response内容
                byte[] bytes = new byte[buffer.readableByteCount()];
                buffer.read(bytes);
                traceData.getHttpResponse().setBody(truncate(new String(bytes, Charset.forName("UTF8"))));
                buffer.readPosition(0);
                setNodeNameWithMappedController(traceData);
            }
        }
        return buffer;
    }

}
